#if defined(_MSC_VER)
#pragma once
#endif

#ifndef LM_CORE_BXDF_H
#define LM_CORE_BXDF_H

#include "lm.h"
#include "spectrum.h"
#include "point.h"
#include "frame.h"
#include "intersection.h"


// Note : TO DO
// Note : How to importance sampling Lafortune BRDF ?


// BxDF utility function
Spectrum FrDiel(float cosi, float cost, const Spectrum &etai, const Spectrum &etat);
Spectrum FrCond(float cosi, const Spectrum &eta, const Spectrum &k);
Spectrum FrApproxEta(const Spectrum &Fr);
Spectrum FrApproxiK(const Spectrum &Fr);


class Fresnel {
public:
	virtual Spectrum Evaluate(float cosi) const = 0;
	~Fresnel() { }
};

class FresnelConductor : public Fresnel {
public:
	FresnelConductor(const Spectrum &e, const Spectrum &kk) : eta(e), k(kk) { }
	Spectrum Evaluate(float cosi) const {
		return FrCond(cosi, eta, k);
	}
	~FresnelConductor() { }
private:
	Spectrum eta, k;
};

class FresnelDielectric : public Fresnel {
public:
	FresnelDielectric(float ei, float et) {
		eta_i = ei;
		eta_t = et;
	}
	~FresnelDielectric() { }
	Spectrum Evaluate(float cosi) const;
private:
	float eta_i, eta_t;
};

class FresnelNoOp : public Fresnel {
public:
	Spectrum Evaluate(float) const { 
		return Spectrum(1.f);
	}
	FresnelNoOp() { }
	~FresnelNoOp() { }
};


enum BxDFType {
	BSDF_REFLECTION			=	1 << 0,
	BSDF_TRANSMISSION		=	1 << 1,
	BSDF_DIFFUSE			=	1 << 2,
	BSDF_GLOSSY				=	1 << 3,
	BSDF_SPECULAR			=	1 << 4,
	BSDF_ALL_TYPES			=	BSDF_DIFFUSE |
								BSDF_GLOSSY	 |
								BSDF_SPECULAR,
	BSDF_ALL_REFLECTION		=	BSDF_REFLECTION |
								BSDF_ALL_TYPES,
	BSDF_ALL_TRANSMISSION	=	BSDF_TRANSMISSION |
								BSDF_ALL_TYPES,
	BSDF_ALL				=	BSDF_ALL_REFLECTION | BSDF_ALL_TRANSMISSION
};

class BxDF {
public:
	BxDF(BxDFType t) : type(t) { }
	~BxDF() { }
	virtual Spectrum f(const Vector3 &wo, const Vector3 &wi) const = 0;
	virtual Spectrum Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &BxDFSample, float *pdf) const;
	virtual Spectrum rho(const Vector3 &wo, int nSamples = 16, Point2 *samples = NULL) const;
	virtual Spectrum rho(int nSamples = 16, Point2 *samples = NULL) const;
	virtual float Pdf(const Vector3 &wo, const Vector3 &wi) const;
	bool MatchesFlags(BxDFType flags) const {
		return (type & flags) == type;
	}
	const BxDFType type;
};

class BRDFToBTDF : public BxDF {
public:
	BRDFToBTDF(BxDF *b)
		: BxDF(BxDFType(b->type ^ (BSDF_REFLECTION | BSDF_TRANSMISSION))) {
			brdf = b;
	}
	~BRDFToBTDF(){ }
	Spectrum f(const Vector3 &wo, const Vector3 &wi) const ;
	Spectrum Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &BxDFSample, float *pdf) const;
	Spectrum rho(const Vector3 &wo, int nSamples = 16, Point2 *samples = NULL) const;
	Spectrum rho(int nSamples = 16, Point2 *samples = NULL) const;
	float Pdf(const Vector3 &wo, const Vector3 &wi) const;
private:
	BxDF *brdf;
};

class ScaledBxDF : public BxDF {
public:
	ScaledBxDF(BxDF *f, const Spectrum sc) : BxDF(BxDFType(f->type)) {
		brdf = f;
		scale = sc;
	}
	~ScaledBxDF(){ }
	Spectrum f(const Vector3 &wo, const Vector3 &wi) const {
		return scale * brdf->f(wo, wi);
	}
	Spectrum rho(const Vector3 &wo, int nSamples = 16, Point2 *samples = NULL) const {
		return scale * brdf->rho(wo, nSamples, samples);
	}
	Spectrum rho(int nSamples = 16, Point2 *samples = NULL) const {
		return scale * brdf->rho(nSamples, samples);
	}
	Spectrum Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &BxDFSample, float *pdf) const {
		Spectrum f = brdf->Sample_f(wo, wi, BxDFSample, pdf);
		return scale * f;
	}
	float Pdf(const Vector3 &wo, const Vector3 &wi) const {
		return brdf->Pdf(wo, wi);
	}
private:
	BxDF *brdf;
	Spectrum scale;
};

class SpecularReflection : public BxDF {
public:
	SpecularReflection(const Spectrum &r, Fresnel *f) : BxDF(BxDFType(BSDF_SPECULAR|BSDF_REFLECTION)) {
		R = r;
		fresnel = f;
	}
	~SpecularReflection() { }
	Spectrum Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &BxDFSample, float *pdf) const;
	Spectrum f(const Vector3 &wo, const Vector3 &wi) const {
		return Spectrum(0.f);
	}
	float Pdf(const Vector3 &wo, const Vector3 &wi) const {
		return 0.0;
	}
private:
	Spectrum R;
	Fresnel *fresnel;
};

class SpecularTransmission : public BxDF {
public:
	SpecularTransmission(const Spectrum &t, float ei, float et) : BxDF(BxDFType(BSDF_SPECULAR|BSDF_TRANSMISSION)), fresnel(ei, et) {
		T = t;
		eta_i = ei;
		eta_t = et;
	}
	~SpecularTransmission() { }
	Spectrum Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &BxDFSample, float *pdf) const;
	Spectrum f(const Vector3 &wo, const Vector3 &wi) const {
		return Spectrum(0.f);
	}
	float Pdf(const Vector3 &wo, const Vector3 &wi) const {
		return 0.0;
	}
private:
	Spectrum T;
	float eta_i, eta_t;
	FresnelDielectric fresnel;
};

class Lambertian : public BxDF {
public:
	Lambertian(const Spectrum &r) : BxDF(BxDFType(BSDF_DIFFUSE | BSDF_REFLECTION)), R(r) { }
	~Lambertian() { }
	Spectrum f(const Vector3 &wo, const Vector3 &wi) const {
		return R * INV_PI;
	}
	Spectrum rho(const Vector3, int, Point2 *) const {
		return R;
	}
	Spectrum rho(int, Point2*) const {
		return R;
	}
private:
	Spectrum R;
};

class OrenNayar : public BxDF {
public:
	OrenNayar(const Spectrum &r, float sig) : BxDF(BxDFType(BSDF_DIFFUSE|BSDF_REFLECTION)), R(r) {
		// sig is degree, turn to radian
		float sigma = Radians(sig);
		float sigma2 = sigma * sigma;
		// wtf...how can u figure out these formula.. 
		A = 1 - (sigma2 / (2*(sigma2+0.33)));
		B = 0.45 * sigma2 / (sigma2 + 0.09);
	}
	~OrenNayar() { }
	Spectrum f(const Vector3 &wo, const Vector3 &wi) const;
private:
	Spectrum R;
	float A, B;
};

class MicrofacetDistribution {
public:
	virtual ~MicrofacetDistribution() { }
	virtual float D(const Vector3 &wh) const = 0;
	virtual void Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &distributionSample, float *pdf) const = 0;
	virtual float Pdf(const Vector3 &wo, const Vector3 &wi) const = 0;
};

class Blinn : public MicrofacetDistribution {
public:
	Blinn(float e) {
		if (e > 10000.f || isnan(e)) e = 10000.f;
		exponent = e;
	}
	~Blinn() { }
	float D(const Vector3 &wh) const {
		float costhetah = fabsf(Frame::CosTheta(wh));
		return (exponent + 2) * INV_TWOPI * powf(costhetah, exponent);
	}
	void Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &distributionSample, float *pdf) const;
	float Pdf(const Vector3 &wo, const Vector3 &wi) const;
private:
	float exponent;
};

class Anisotropic : public MicrofacetDistribution {
public:
	Anisotropic(float x, float y) {
		ex = x;
		ey = y;
		if (ex > 10000.f || isnan(ex)) ex = 10000.f;
		if (ey > 10000.f || isnan(ey)) ey = 10000.f;
	}
	~Anisotropic() { }
	float D(const Vector3 &wh) const {
		float costhetah = Frame::AbsCosTheta(wh);
		float d = 1.f - costhetah * costhetah;
		if (d == 0.f) return 0.f;
		float e = (ex * wh.x * wh.x + ey * wh.y * wh.y) / d;
		return sqrtf((ex+2.f) * (ey+2.f)) * INV_TWOPI * powf(costhetah, e);
	}
	void Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &distributionSample, float *pdf) const;
	void SampleFirstQuadrant(const Point2 &sample, float *phi, float *costheta) const;
	float Pdf(const Vector3 &wo, const Vector3 &wi) const;
private:
	float ex, ey;
};

class Microfacet : public BxDF {
public:
	Microfacet(const Spectrum &r, Fresnel *f, MicrofacetDistribution *d) : 
	  BxDF(BxDFType(BSDF_REFLECTION | BSDF_GLOSSY)), R(r), distribution(d), fresnel(f) { }
	~Microfacet() { }
	Spectrum f(const Vector3 &wo, const Vector3 &wi) const;
	float G(const Vector3 &wo, const Vector3 &wi, const Vector3 &wh) const {
		float NdotWh = Frame::AbsCosTheta(wh);
		float NdotWo = Frame::AbsCosTheta(wo);
		float NdotWi = Frame::AbsCosTheta(wi);
		float WOdotWh = AbsDot(wo, wh);
		return min(1.f, min((2.f * NdotWh * NdotWo / WOdotWh),
			(2.f * NdotWh * NdotWi / WOdotWh)));
	}
	Spectrum Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &bxdfSample, float *pdf) const;
	float Pdf(const Vector3 &wo, const Vector3 &wi) const;
private:
	Spectrum R;
	MicrofacetDistribution *distribution;
	Fresnel *fresnel;
};

class Lafortune : public BxDF {
public:
	Lafortune(const Spectrum &r, u_int nl, const Spectrum *xx, const Spectrum *yy, const Spectrum *zz, const Spectrum *e, BxDFType type)
		: BxDF(type), R(r) {
		x = xx;
		y = yy;
		z = zz;
		exponent = e;
	}
	~Lafortune() { }
	Spectrum f(const Vector3 &wo, const Vector3 &wi) const;
private:
	Spectrum R;
	u_int nLobes;
	const Spectrum *x, *y, *z, *exponent;
};

class FresnelBlend : public BxDF {
public:
	FresnelBlend(const Spectrum &rd, const Spectrum &rs, MicrofacetDistribution *dist)
		: BxDF(BxDFType(BSDF_GLOSSY|BSDF_REFLECTION)), Rd(rd), Rs(rs) {
			distribution = dist;
	}
	~FresnelBlend() { }
	Spectrum f(const Vector3 &wo, const Vector3 &wi) const;
	Spectrum SchilckFresnel(float costheta) const {
		return Rs + powf(1-costheta, 5.f) * (Spectrum(1.0f) - Rs);
	}
	Spectrum Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &bxdfSample, float *pdf) const;
	float Pdf(const Vector3 &wo, const Vector3 &wi) const;
private:
	Spectrum Rd, Rs;
	MicrofacetDistribution *distribution;
};

// for ward, we only use the specular term
// when used, we can combine it with a Lambertian term
// only need to make sure that Rs + Rd  <= 1
class Ward : public BxDF {
public:
	Ward(const Spectrum rs, float ax, float ay) : BxDF(BxDFType(BSDF_REFLECTION|BSDF_GLOSSY)), Rs(rs), alphaX(ax), alphaY(ay) {
		INV_FOURPI_ALPHAXY = 1.f / alphaX / alphaY * INV_FOURPI;
		invAlphaX2 = 1.f / (alphaX*alphaX);
		invAlphaY2 = 1.f / (alphaY*alphaY);
	}
	~Ward() { }
	Spectrum f(const Vector3 &wo, const Vector3 &wi) const;
	Spectrum Sample_f(const Vector3 &wo, Vector3 *wi, const Point2 &BxDFSample, float *pdf) const;
	float Pdf(const Vector3 &wo, const Vector3 &wi) const;
private:
	Spectrum Rs;
	float alphaX, alphaY;
	float invAlphaX2, invAlphaY2;
	float INV_FOURPI_ALPHAXY;
};

#endif	// LM_CORE_BXDF_H